import argparse
import json
import re
import numpy
import numpy as np

from src.util import GraphUtil
from src.util.GraphUtil import BarGraphUtil, ScatterGraphUtil


class ArgumentParser:
    def __init__(self):
        self.parser = argparse.ArgumentParser()
        self.parser.add_argument("-f", "--files", type=str, help="source file(s) path to generate graph", nargs="+",
                                 required=True)
        self.parser.add_argument("-x", "--xlabel", type=str, help="name of xlable", default="X")
        self.parser.add_argument("-y", "--ylabel", type=str, help="name of ylable", default="Y")
        self.parser.add_argument("-n", "--names", type=str, help="name(s) of models", nargs="+", required=True)
        self.parser.add_argument("-t", "--type", type=str, help="type of graph", choices=["bar", "scatter"],
                                 required=True)
        self.parser.add_argument("-o", "--output", type=str, help="output file",
                                 default="plot")

    def parse_args(self):
        args = self.parser.parse_args()
        files = args.files
        xlabel = args.xlabel
        ylabel = args.ylabel
        type = args.type
        names = args.names
        assert len(files) == len(names), "length of source files and model names must be same"
        output_file = args.output
        return files, xlabel, ylabel, type, names, output_file


def parse_score(text, type):
    pattern = re.compile(rf'{re.escape(type)}:\s*(\d*\.\d+|\d+(\.\d*)?)')
    match = pattern.search(text)
    if match:
        return float(match.group(1))
    else:
        return None


if __name__ == "__main__":
    arg_parser = ArgumentParser()
    files, xlabel, ylabel, type, names, output_file = arg_parser.parse_args()
    graphUtil = GraphUtil()
    if type == "bar":
        graphUtil = BarGraphUtil()
        datas = []
        for file in files:
            data_F1 = []
            data_P = []
            data_R = []
            with open(file, 'r', encoding='utf-8') as f:
                for line in f:
                    data = json.loads(line)
                    score = data['score']
                    data_F1.append(float(parse_score(score, 'F1')))
                    data_P.append(float(parse_score(score, 'P')))
                    data_R.append(float(parse_score(score, 'R')))
                f.close()
            avg = []
            avg.append(np.mean(data_P))
            avg.append(np.mean(data_R))
            avg.append(np.mean(data_F1))
            datas.append(avg)
        graphUtil.setData(datas, xlabel, ylabel, names)
        graphUtil.save(output_file + '.png')
    elif type == 'scatter':
        graphUtil = ScatterGraphUtil()
        datas_F1 = []
        datas_P = []
        datas_R = []
        for file in files:
            data_F1 = []
            data_P = []
            data_R = []
            with open(file, 'r', encoding='utf-8') as f:
                for line in f:
                    data = json.loads(line)
                    score = data['score']
                    data_F1.append(float(parse_score(score, 'F1')))
                    data_P.append(float(parse_score(score, 'P')))
                    data_R.append(float(parse_score(score, 'R')))
                f.close()
            if len(datas_F1) != 0:
                assert len(data_F1) == len(datas_F1[0]), 'different length of files '
            datas_F1.append(data_F1)
            datas_P.append(data_P)
            datas_R.append(data_R)
        datas = [datas_F1, datas_P, datas_R]
        flag = ['F1', 'P', 'R']
        for index in range(0, 3):
            graphUtil.setData(datas[index], xlabel, ylabel, names)
            graphUtil.setTitle(flag[index] + '(' + str(len(datas[index][0])) + ')')
            graphUtil.save(output_file + '_' + flag[index] + '.png')
